import numpy as np
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import math
import random
import os
import time

class PSO:

    def __init__(self, D, N, M, p_low, p_up, v_low, v_high,ax, w = 1., c1 = 0.5, c2 = 0.5):
        self.w = w  # 惯性权值
        self.c1 = c1  # 个体学习因子
        self.c2 = c2  # 群体学习因子
        self.D = D  # 粒子维度
        self.N = N  # 粒子群规模，初始化种群个数
        self.M = M  # 最大迭代次数
        self.p_range = [p_low, p_up]  # 粒子位置的约束范围
        self.v_range = [v_low, v_high]  # 粒子速度的约束范围
        self.x = np.zeros((self.N, self.D))  # 所有粒子的位置
        self.v = np.zeros((self.N, self.D))  # 所有粒子的速度
        self.p_best = np.zeros((self.N, self.D))  # 每个粒子的最优位置
        self.g_best = np.zeros((1, self.D))[0]  # 种群（全局）的最优位置
        self.p_bestFit = np.zeros(self.N)  # 每个粒子的最优适应值
        self.g_bestFit = float('Inf')  # float('-Inf')，始化种群（全局）的最优适应值，由于求极小值，故初始值给大，向下收敛，这里默认优化问题中只有一个全局最优解
        self.ax = ax

        # 初始化所有个体和全局信息
        for i in range(self.N):
            for j in range(self.D):
                self.x[i][j] = random.uniform(self.p_range[0][j], self.p_range[1][j])
                self.v[i][j] = random.uniform(self.v_range[0], self.v_range[1])
            self.p_best[i] = self.x[i]  # 保存个体历史最优位置，初始默认第0代为最优
            fit = self.fitness(self.p_best[i])
            self.p_bestFit[i] = fit  # 保存个体历史最优适应值
            if fit < self.g_bestFit:  # 寻找并保存全局最优位置和适应值
                self.g_best = self.p_best[i]
                self.g_bestFit = fit
        
        #初始化标点
        chat3dplot(self.x,self.ax,self.fitness,'o','black')

    def fitness(self,x):
        func_value = 0
        for i in range(len(x)):
            func_value += x[i]**2
        return func_value

    def update(self):
        for i in range(self.N):
            self.v[i] = self.w * self.v[i] + self.c1 * random.uniform(0, 1) * (self.p_best[i] - self.x[i]) + self.c2 * random.uniform(0, 1) * (self.g_best - self.x[i])
            #速度限制
            for j in range(self.D):
                if self.v[i][j] < self.v_range[0]:
                    self.v[i][j] = self.v_range[0]
                if self.v[i][j] > self.v_range[1]:
                    self.v[i][j] = self.v_range[1]
            #更新位置
            self.x[i] = self.x[i] + self.v[i]
            #位置限制
            for j in range(self.D):
                if self.x[i][j] < self.p_range[0][j]:
                    self.x[i][j] = self.p_range[0][j]
                if self.x[i][j] > self.p_range[1][j]:
                    self.x[i][j] = self.p_range[1][j]
            #更新个体和全局历史最优位置及适应值
            _fit = self.fitness(self.x[i])
            # print(self.x[i])
            if _fit < self.p_bestFit[i]:
                self.p_best[i] = self.x[i]
                self.p_bestFit[i] = _fit
            if _fit < self.g_bestFit:
                self.g_best = self.x[i]
                self.g_bestFit = _fit
    
    def pso(self,draw = 1):
        best_fit = [] #记录每轮迭代最佳适应度，用于绘图
        w_range = None
        #plt.show()
        if isinstance(self.w,tuple):
            w_range = self.w[1] - self.w[0]
            self.w = self.w[1]
        time_start = time.time() #记录迭代寻优开始时间
        for i in range(self.M):
            self.update()
            if w_range:
                self.w -= w_range / self.M
            print("\rIter：{:d}/{:d} fitness:{:.4f}".format(i,self.M,self.g_bestFit,end = '\n'))
            best_fit.append(self.g_bestFit.copy())
        time_end = time.time() #记录迭代结束时间
        chat3dplot(self.x,self.ax,self.fitness,'+','red')#3d图标点
        print(f'Algorithm takes {time_end - time_start} seconds')
        if draw:
            plt.figure()
            plt.semilogy([i for i in range(self.M)],best_fit)
            plt.xlabel("iter")
            plt.ylabel("fitness")
            plt.title("Iter process")
            plt.show()
            

#给三维图标点
def chat3dplot(x,ax,fitness,xz,color):
    a=[]
    b=[]
    c=[]
    for i in range(len(x)):
        a.append(x[i][0])
        b.append(x[i][1])
        c.append(fitness(x[i]))

    ax.plot(a,b,c,c = color, marker = xz,zorder=10,linestyle='')    

    

if __name__ == '__main__':
    low = [-5.12,-5.12]
    up = [5.12,5.12]
    #绘制三维函数图
    fig = plt.figure()
    ax = fig.add_axes(Axes3D(fig))
    x = np.arange(-5.12,5.12,0.05)
    y = np.arange(-5.12,5.12,0.05)
    x,y = np.meshgrid(x,y)
    z = x**2 + y**2
        

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    ax.plot_surface(x, y, z, cmap='rainbow',zorder=1)
    
    pso = PSO(2, 10, 100, low, up, -0.5, 0.5, ax,w = 0.8)
    pso.pso()

